{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "''' notebook needed to reproduce MSK survival analysis\n",
    "    (fig 2D/E) '''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msk = pd.read_csv('../Data_input/msk_impact_2017/MSK-IMPACT_cosmic_tier1.txt', sep='\\t')\n",
    "msk = msk.drop(['GENE_ID'], axis=1)\n",
    "msk = msk.set_index('COMMON')\n",
    "msk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "patient_meta = pd.read_csv('../Data_input/msk_impact_2017/data_clinical_patient_edit.txt', sep='\\t')\n",
    "patient_meta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define list of patients we have in our msk cosmic tier1 df\n",
    "my_patients = list(msk.columns)\n",
    "my_patients.remove('Unnamed: 10947')\n",
    "\n",
    "my_patients_trimmed = [x.split('-')[0] + '-' + x.split('-')[1] for x in my_patients] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# didnt actually remove anything, huh\n",
    "patient_meta = patient_meta[patient_meta.PATIENT_ID.isin(my_patients_trimmed)]\n",
    "patient_meta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#//////////////////////////////////////////////////////////////////////\n",
    "#//////////////////////////////////////////////////////////////////////\n",
    "#/////////////    what does our data look like?      //////////////////\n",
    "#//////////////////////////////////////////////////////////////////////"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cov = pd.read_csv('../Data_input/mutation_input/coverage_all_cells_cerebra.csv')\n",
    "cov"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cov.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "list(cov.columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#/////////////////////////////////////////////////////////////////\n",
    "#/////////////////////////////////////////////////////////////////\n",
    "#/////////////////////////////////////////////////////////////////"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get list of all the unique mutations present in the msk dataset\n",
    "msk_muts = []\n",
    "\n",
    "for idx, row in msk.iterrows():\n",
    "    l = list(row)\n",
    "\n",
    "    for elm in l:\n",
    "        \n",
    "        if not pd.isna(elm):\n",
    "            if ',' in elm:\n",
    "                elm_split = elm.split(',')\n",
    "                \n",
    "                for sub_elm in elm_split:\n",
    "                    mut = idx + '_' + sub_elm\n",
    "                    msk_muts.append(mut)\n",
    "            else:\n",
    "                mut = idx + '_' + elm\n",
    "                msk_muts.append(mut)\n",
    "                \n",
    "msk_muts_unique = list(set(msk_muts))\n",
    "msk_muts_unique"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs_muts = list(cov.columns)\n",
    "obs_muts.remove('sample')\n",
    "obs_muts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(msk_muts_unique))\n",
    "print(len(obs_muts))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#/////////////////////////////////////////////////////////////////\n",
    "#/////////////////////////////////////////////////////////////////\n",
    "#// need to convert msk df into a by-ROI df that resembles ours //\n",
    "#/////////////////////////////////////////////////////////////////\n",
    "#/////////////////////////////////////////////////////////////////"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make a new df, oriented just like ours\n",
    "msk_reorient = pd.DataFrame(index=msk_muts_unique, columns=msk.columns)\n",
    "msk_reorient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msk_reorient.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# only 141 muts *perfectly* conserved across both\n",
    "msk_reorient_filt = msk_reorient[msk_reorient.index.isin(obs_muts)]\n",
    "msk_reorient_filt "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# whats in obs_muts but NOT IN msk_muts_unique? \n",
    "#   this is worrisome. should come back to this later\n",
    "set(obs_muts) - set(msk_muts_unique)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# init new msk df values to zero\n",
    "msk_reorient_filt[:] = 0\n",
    "#msk_reorient_filt\n",
    "msk_reorient_filt = msk_reorient_filt.drop(['Unnamed: 10947'], axis=1)\n",
    "msk_reorient_filt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fill in msk_reorient_filt\n",
    "for idx, row in msk.iterrows():\n",
    "    for i in range(0,len(row.index)):\n",
    "        elm = row[i]\n",
    "        patient_id = row.index[i]\n",
    "\n",
    "        if not pd.isna(elm):\n",
    "            \n",
    "            if ',' in elm:\n",
    "                elm_split = elm.split(',')\n",
    "                for sub_elm in elm_split:\n",
    "                    mut_str = idx + '_' + sub_elm\n",
    "            else:    \n",
    "                mut_str = idx + '_' + elm\n",
    "            \n",
    "            if mut_str in list(msk_reorient_filt.index):\n",
    "                msk_reorient_filt.loc[mut_str, patient_id] += 1 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# validate\n",
    "    # works! tho df looks like all zeros\n",
    "msk_reorient_filt.loc['KRAS_G12A', 'P-0002921-T01-IM3']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msk_reorient_filt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#/////////////////////////////////////////////////////////////////\n",
    "#/////////////////////////////////////////////////////////////////\n",
    "#//////    now lets get cov to look exactly the same way   ///////\n",
    "#/////////////////////////////////////////////////////////////////"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cov = cov.set_index('sample')\n",
    "cov"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cov_binary = cov.applymap(lambda x: int(x.split(':')[0]))\n",
    "cov_binary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def binarize(x):\n",
    "    \"\"\" binarize a dataframe; 1 for any non-zero value \"\"\"\n",
    "    if x != 0:\n",
    "        return(1)\n",
    "    else:\n",
    "        return(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cov_binary = cov_binary.applymap(binarize)\n",
    "cov_binary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# transpose, so orientation matches msk\n",
    "cov_binary_T = cov_binary.transpose()\n",
    "cov_binary_T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# subset muts, so that they match msk_reorient_filt\n",
    "cov_binary_T = cov_binary_T[cov_binary_T.index.isin(list(msk_reorient_filt.index))]\n",
    "cov_binary_T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#/////////////////////////////////////////////////////////////////\n",
    "#/////////////////////////////////////////////////////////////////\n",
    "#////////    want to break this thing down by SAMPLE    //////////\n",
    "#/////////////////////////////////////////////////////////////////"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "meta = pd.read_csv('../metadata_input/metadata_all_cells.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create list of all the sample we have cells from \n",
    "all_samples = []\n",
    "\n",
    "for cell in list(cov_binary_T.columns):\n",
    "    meta_line = meta[meta.cell_id == cell]\n",
    "    sample = list(meta_line.sample_name)[0]\n",
    "    all_samples.append(sample)\n",
    "    \n",
    "all_samples_u = list(set(all_samples))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create new, by-sample dataframe \n",
    "muts_by_sample_df = pd.DataFrame(index = cov_binary_T.index, columns=all_samples_u)\n",
    "muts_by_sample_df[:] = 0\n",
    "muts_by_sample_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fill in by-sample df \n",
    "for idx, row in cov_binary_T.iterrows():\n",
    "    for i in range(0,len(row)):\n",
    "        curr_cell = row.index[i]\n",
    "        curr_val = row[i]\n",
    "        \n",
    "        meta_line = meta[meta.cell_id == curr_cell]\n",
    "        curr_sample = list(meta_line.sample_name)[0]\n",
    "        \n",
    "        if curr_val == 1 and muts_by_sample_df.loc[idx,curr_sample] == 0:\n",
    "            muts_by_sample_df.loc[idx,curr_sample] = 1\n",
    "            \n",
    "muts_by_sample_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msk_reorient_filt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# write to csv\n",
    "#muts_by_sample_df.to_csv('muts_by_sample.csv')\n",
    "#msk_reorient_filt.to_csv('msk_reorient.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get msk CLINICAL meta, and subset by LAUD patients\n",
    "msk_clinical_meta = pd.read_csv(\n",
    "    '../Data_input/msk_impact_2017/data_clinical_sample.txt', sep='\\t')\n",
    "\n",
    "msk_clinical_meta_laud = msk_clinical_meta[msk_clinical_meta.CANCER_TYPE_DETAILED == 'Lung Adenocarcinoma']\n",
    "msk_clinical_meta_laud = msk_clinical_meta_laud.reset_index(drop=True)\n",
    "\n",
    "# remove duplicated patients\n",
    "msk_clinical_meta_laud = msk_clinical_meta_laud.loc[~msk_clinical_meta_laud.index.duplicated(keep='first')]\n",
    "\n",
    "display(msk_clinical_meta_laud.shape)\n",
    "\n",
    "msk_laud_patients = list(msk_clinical_meta_laud.PATIENT_ID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read in msk meta SURVIVAL data\n",
    "path = '../Data_input/msk_impact_2017/data_clinical_patient.txt'\n",
    "meta_survival = pd.read_csv(path, sep='\\t')\n",
    "meta_survival_sub = meta_survival[meta_survival.OS_MONTHS.notna()] # subset by those patients that HAVE survival data\n",
    "meta_survival_sub = meta_survival_sub.reset_index(drop=True)\n",
    "meta_survival_sub['OS_YEARS'] = meta_survival_sub.OS_MONTHS / 12 # adding in a OS (years) col\n",
    "print(meta_survival_sub.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# remove sample IDs from labels\n",
    "msk_reindex = [x.split('-')[0] + '-' + x.split('-')[1] for x in list(msk.index)]\n",
    "msk.index = msk_reindex\n",
    "\n",
    "# remove duplicated patients\n",
    "msk = msk.loc[~msk.index.duplicated(keep='first')]\n",
    "\n",
    "# subset by LAUD patients\n",
    "msk = msk[msk.index.isin(msk_laud_patients)]\n",
    "print(msk.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the samples that have EGFR L858R muts\n",
    "sub_col = msk[['EGFR_L858R']]\n",
    "nonzero_arr = sub_col.EGFR_L858R.nonzero()\n",
    "egfr_l858r_samples = list(sub_col.index[nonzero_arr])\n",
    "\n",
    "print(len(egfr_l858r_samples))\n",
    "#egfr_l858r_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get all the other samples\n",
    "non_egfr_l858r_samples = list(msk.loc[~msk.index.isin(egfr_l858r_samples)].index)\n",
    "print(len(non_egfr_l858r_samples))\n",
    "#non_egfr_l858r_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#///////////////////////////////////////////////////////////////////////////////////\n",
    "#///////////////////////////////////////////////////////////////////////////////////\n",
    "#//////////    what if we take a look at all of our 'drivers',     /////////////////\n",
    "#//////////         not just the EGFR guys          ////////////////////////////////\n",
    "#/////////////////////////////////////////////////////////////////////////////////// "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fuck it \n",
    "    # lets do all the drivers at the same time\n",
    "drivers_list = ['KRAS_G12C', 'KRAS_Q61K', 'KRAS_C185S', 'KRAS_Q61H', 'KRAS_G13C', 'KRAS_G13D', 'KRAS_G12V', \n",
    "                'KRAS_G12F', 'KRAS_G12S', 'KRAS_G12C', 'KRAS_G12A', 'BRAF_V600E', 'EGFR_L858R', 'EGFR_del19']\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# how many of these are actually represented in our dataset? \n",
    "    # not many...\n",
    "    # the lack of EGFR_del19 is concerning\n",
    "    \n",
    "found_muts_list = []\n",
    "for x in list(msk.columns):\n",
    "    if x in drivers_list:\n",
    "        found_muts_list.append(x)\n",
    "        \n",
    "print(found_muts_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the samples that have 'driver' mutations\n",
    "msk_drivers_club = msk[found_muts_list]\n",
    "msk_drivers_club['rowsum'] = msk_drivers_club.sum(axis=1)\n",
    "\n",
    "msk_drivers_club = msk_drivers_club[msk_drivers_club.rowsum > 0]\n",
    "msk_drivers_club"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "driver_mut_patients = list(msk_drivers_club.index)\n",
    "print(len(driver_mut_patients))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# interesting that none of these guys have > 1 mutation \n",
    "    # ohhh heyyyYy maybe single cell is valuable after all?? \n",
    "set(msk_drivers_club.rowsum)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get all the other samples\n",
    "non_driver_mut_patients = list(msk.loc[~msk.index.isin(driver_mut_patients)].index)\n",
    "print(len(non_driver_mut_patients))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#///////////////////////////////////////////////////////////////////////////////////\n",
    "#///////////////////////////////////////////////////////////////////////////////////\n",
    "#//////    what if i dont restrict myself to the 141 'shared' mutations?     ///////\n",
    "#///////////////////////////////////////////////////////////////////////////////////"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msk_raw = pd.read_csv('../Data_input/msk_impact_2017/MSK-IMPACT_cosmic_tier1.txt', sep='\\t')\n",
    "msk_raw = msk_raw.drop(['GENE_ID'], axis=1)\n",
    "msk_raw = msk_raw.set_index('COMMON')\n",
    "#msk_raw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# need to REFORMAT into something were more familiar with\n",
    "    # FIRST get list of all the unique mutations present in the msk dataset\n",
    "msk_raw_muts = []\n",
    "\n",
    "for idx, row in msk_raw.iterrows():\n",
    "    l = list(row)\n",
    "\n",
    "    for elm in l:\n",
    "        if not pd.isna(elm):\n",
    "            if ',' in elm:\n",
    "                elm_split = elm.split(',')\n",
    "                \n",
    "                for sub_elm in elm_split:\n",
    "                    mut = idx + '_' + sub_elm\n",
    "                    msk_raw_muts.append(mut)\n",
    "            else:\n",
    "                mut = idx + '_' + elm\n",
    "                msk_raw_muts.append(mut)\n",
    "                \n",
    "msk_raw_muts_unique = list(set(msk_raw_muts))\n",
    "msk_raw_muts_unique"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# trim sample names\n",
    "samples_list_orig = list(msk_raw.columns)\n",
    "samples_list_orig.remove('Unnamed: 10947')\n",
    "samples_list_trimmed = [x.split('-')[0] + '-' + x.split('-')[1] for x in samples_list_orig]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make a new df, and init values to zero\n",
    "msk_raw_reorient = pd.DataFrame(index=samples_list_trimmed, columns=msk_raw_muts_unique)\n",
    "msk_raw_reorient[:] = 0\n",
    "#msk_raw_reorient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fill in msk_raw_reorient\n",
    "for idx, row in msk_raw.T.iterrows():\n",
    "    try:\n",
    "        sample_name_trimmed = idx.split('-')[0] + '-' + idx.split('-')[1]\n",
    "    except IndexError:\n",
    "        print(sample_name_trimmed)\n",
    "        \n",
    "    for i in range(0,len(row)):\n",
    "        gene = row.index[i]\n",
    "        mut = row.iloc[i]\n",
    "    \n",
    "        if not pd.isna(mut):\n",
    "            mut_str = gene + '_' + mut\n",
    "            \n",
    "            # now update msk_raw_reorient\n",
    "            if mut_str in list(msk_raw_reorient.columns):\n",
    "                msk_raw_reorient.loc[sample_name_trimmed, mut_str] += 1 # for some reason this \n",
    "                #msk_raw_reorient.loc[sample_name_trimmed, mut_str] = 1    # is like 10x faster?? truly dont understand whats going on here "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# binarize the mutations matrix -- dont want values > 1\n",
    "def binarize(val):\n",
    "    if val >= 1:\n",
    "        return(1)\n",
    "    else:\n",
    "        return(0)\n",
    "\n",
    "msk_raw_reorient_b = msk_raw_reorient.applymap(binarize)\n",
    "msk_raw_reorient_b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sanity check -- did binarize() actually do anything? \n",
    "    # should return 0 and 1 \n",
    "np.unique(msk_raw_reorient_b.values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# lets subset down to just LAUD patients\n",
    "msk_raw_reorient_b = msk_raw_reorient_b[msk_raw_reorient_b.index.isin(msk_laud_patients)]\n",
    "msk_raw_reorient_b.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# what EGFR muts are we actually seeing? \n",
    "EGFR_del19_flavors = []\n",
    "\n",
    "for x in list(msk_raw_reorient_b.columns):\n",
    "    if 'EGFR' in x:\n",
    "        if 'del' in x or 'ins' in x:\n",
    "            EGFR_del19_flavors.append(x)\n",
    "            \n",
    "EGFR_del19_flavors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# now lets see how many of these other drivers we can find\n",
    "drivers_list = ['KRAS_G12C', 'KRAS_Q61K', 'KRAS_C185S', 'KRAS_Q61H', 'KRAS_G13C', 'KRAS_G13D', 'KRAS_G12V', \n",
    "                'KRAS_G12F', 'KRAS_G12S', 'KRAS_G12A', 'BRAF_V600E', 'EGFR_L858R']\n",
    "\n",
    "found_muts = []\n",
    "for mut_str in drivers_list:\n",
    "    if mut_str in list(msk_raw_reorient_b.columns):\n",
    "        found_muts.append(mut_str)\n",
    "\n",
    "print(found_muts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# subset msk_raw dataframe\n",
    "all_found_muts = found_muts + EGFR_del19_flavors\n",
    "\n",
    "msk_raw_drivers_subset = msk_raw_reorient_b[all_found_muts]\n",
    "msk_raw_drivers_subset['rowsum'] = msk_raw_drivers_subset.sum(axis=1)\n",
    "\n",
    "# include only samples with > 0 mutations\n",
    "msk_raw_drivers_subset = msk_raw_drivers_subset[msk_raw_drivers_subset.rowsum > 0]\n",
    "msk_raw_drivers_subset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get samples that have driver muts\n",
    "driver_mut_patients = list(msk_raw_drivers_subset.index)\n",
    "\n",
    "# get all the other samples\n",
    "non_driver_mut_patients = list(msk_raw_reorient_b.loc[~msk_raw_reorient_b.index.isin(driver_mut_patients)].index)\n",
    "\n",
    "driver_mut_patients = list(set(driver_mut_patients))\n",
    "non_driver_mut_patients = list(set(non_driver_mut_patients))\n",
    "\n",
    "print(len(driver_mut_patients))\n",
    "print(len(non_driver_mut_patients))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#///////////////////////////////////////////////////////////////////////////////////\n",
    "#///////////////////////////////////////////////////////////////////////////////////\n",
    "#//////    now that we have the patients lists, lets look at survival     //////////\n",
    "#///////////////////////////////////////////////////////////////////////////////////"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# subset survival metadata frames by in/out group patients lists\n",
    "driver_patients_survival_meta = meta_survival_sub[meta_survival_sub.PATIENT_ID.isin(driver_mut_patients)]\n",
    "driver_patients_survival_meta = driver_patients_survival_meta.reset_index(drop=True)\n",
    "\n",
    "non_driver_patients_survival_meta = meta_survival_sub[meta_survival_sub.PATIENT_ID.isin(non_driver_mut_patients)]\n",
    "non_driver_patients_survival_meta = non_driver_patients_survival_meta.reset_index(drop=True)\n",
    "\n",
    "print(driver_patients_survival_meta.shape)\n",
    "print(non_driver_patients_survival_meta.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# lets start by just looking at some distributions\n",
    "squish = pd.DataFrame(driver_patients_survival_meta.OS_MONTHS)\n",
    "squish = squish.rename(columns={'OS_MONTHS': 'SURVIVAL_DRIVER'})\n",
    "squish['SURVIVAL_nonDRIVER'] = non_driver_patients_survival_meta.OS_MONTHS\n",
    "\n",
    "# looks pretty good when i supress the outliers\n",
    "squish.boxplot(column=['SURVIVAL_DRIVER', 'SURVIVAL_nonDRIVER'], showfliers=False, showmeans=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# maybe i should just be looking at the deceased guys\n",
    "driver_patients_survival_meta_d = driver_patients_survival_meta[driver_patients_survival_meta.OS_STATUS == 'DECEASED']\n",
    "non_driver_patients_survival_meta_d = non_driver_patients_survival_meta[non_driver_patients_survival_meta.OS_STATUS == 'DECEASED']\n",
    "\n",
    "squish_d = pd.DataFrame(driver_patients_survival_meta_d.OS_MONTHS)\n",
    "squish_d = squish_d.rename(columns={'OS_MONTHS': 'SURVIVAL_DRIVER'})\n",
    "squish_d['SURVIVAL_nonDRIVER'] = non_driver_patients_survival_meta_d.OS_MONTHS\n",
    "\n",
    "# looks pretty good when i supress the outliers\n",
    "squish_d.boxplot(column=['SURVIVAL_DRIVER', 'SURVIVAL_nonDRIVER'], showfliers=False, showmeans=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import ttest_ind\n",
    "t, p = ttest_ind(squish.SURVIVAL_DRIVER, squish.SURVIVAL_nonDRIVER, equal_var=False)\n",
    "print(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def kmf_setup_refined(meta_, in_samples, out_samples):\n",
    "    \"\"\" can we set this up so were passing in only metadata and in/out lists? \"\"\"\n",
    "    meta_in = meta_[meta_.PATIENT_ID.isin(in_samples)]\n",
    "    meta_out = meta_[meta_.PATIENT_ID.isin(out_samples)]\n",
    "    \n",
    "    # get OS_MONTHS/YEARS and OS_STATUS\n",
    "    dur_in = list(meta_in.OS_YEARS)\n",
    "    event_in = [x == 'DECEASED' for x in list(meta_in.OS_STATUS)]\n",
    "    \n",
    "    dur_out = list(meta_out.OS_YEARS)\n",
    "    event_out = [x == 'DECEASED' for x in list(meta_out.OS_STATUS)]\n",
    "    \n",
    "    return(dur_in, event_in, dur_out, event_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#///////////////////////////////////////////////////////////////////////////////////\n",
    "#///////////////////////////////////////////////////////////////////////////////////\n",
    "#//////////    break down to just EGFR mut patient set                   ///////////\n",
    "#//////////        makes sense to focus on this bc we dont know what     ///////////\n",
    "#//////////        everyone else has / has been treated with. best       ///////////\n",
    "#//////////        to compare apples to apples                           ///////////\n",
    "#///////////////////////////////////////////////////////////////////////////////////"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# subset raw dataframe\n",
    "egfr_muts_list = EGFR_del19_flavors + ['EGFR_L858R']\n",
    "\n",
    "msk_raw_egfr_subset = msk_raw_reorient_b[egfr_muts_list]\n",
    "msk_raw_egfr_subset['rowsum'] = msk_raw_egfr_subset.sum(axis=1)\n",
    "\n",
    "# include only samples with > 0 mutations\n",
    "msk_raw_egfr_subset = msk_raw_egfr_subset[msk_raw_egfr_subset.rowsum > 0]\n",
    "msk_raw_egfr_subset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msk_egfr_laud_patients = list(set(msk_raw_egfr_subset.index))\n",
    "print(len(msk_egfr_laud_patients))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#///////////////////////////////////////////////////////////////////////////////////\n",
    "#///////////////////////////////////////////////////////////////////////////////////\n",
    "#//////     lets try another method of finding patients with EGFR     //////////////\n",
    "#//////        and non-EGFR muts                                      //////////////\n",
    "#///////////////////////////////////////////////////////////////////////////////////"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# subset to include only our driver muts\n",
    "msk_reorient_drivers_sub = msk_raw_reorient_b[all_found_muts]\n",
    "\n",
    "# get row sums\n",
    "msk_reorient_drivers_sub['rowsum'] = msk_reorient_drivers_sub.sum(axis=1)\n",
    "set(msk_reorient_drivers_sub.rowsum)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msk_reorient_drivers_sub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# just making sure there are other mutations in this damn thing\n",
    "for idx, row in msk_reorient_drivers_sub.iterrows():\n",
    "    for i in range(0,len(row)):\n",
    "        curr_val = row.iloc[i]\n",
    "        curr_mut = row.index[i]\n",
    "    \n",
    "        if curr_val != 0 and 'EGFR' not in curr_mut and curr_mut != 'rowsum':\n",
    "            print(curr_mut)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# just another way of doing the EGFR + non-EGFR patient search\n",
    "    # comes up with nothing, again \n",
    "for idx, row in msk_reorient_drivers_sub.iterrows():\n",
    "    non_egfr_found = 0\n",
    "    egfr_found = 0 \n",
    "    \n",
    "    for i in range(0,len(row)):\n",
    "        curr_val = row.iloc[i]\n",
    "        curr_mut = row.index[i]\n",
    "    \n",
    "        if curr_val != 0 and curr_mut != 'rowsum':\n",
    "            if 'EGFR' not in curr_mut:\n",
    "                non_egfr_found = 1\n",
    "            elif 'EGFR' in curr_mut:\n",
    "                egfr_found = 1\n",
    "                \n",
    "    if non_egfr_found and egfr_found:\n",
    "        print(idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#///////////////////////////////////////////////////////////////////////////////////\n",
    "#///////////////////////////////////////////////////////////////////////////////////\n",
    "#//////   want to add up 'total mutational burden' for our LAUD       //////////////\n",
    "#//////       driver subset                                           //////////////\n",
    "#///////////////////////////////////////////////////////////////////////////////////"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msk_raw_drivers_subset = msk_reorient_drivers_sub.drop('rowsum', axis=1)\n",
    "msk_raw_drivers_subset['EGFR_muts'] = 0\n",
    "msk_raw_drivers_subset['nonEGFR_muts'] = 0\n",
    "msk_raw_drivers_subset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# record frequency of both EGFR and nonEGFR mutations\n",
    "for idx, row in msk_raw_drivers_subset.iterrows():\n",
    "    non_egfr_count = 0\n",
    "    egfr_count = 0\n",
    "\n",
    "    for i in range(0,len(row)):\n",
    "        curr_val = row.iloc[i]\n",
    "        curr_mut = row.index[i]\n",
    "    \n",
    "        if curr_val != 0 and curr_mut != 'EGFR_muts' and curr_mut != 'nonEGFR_muts':\n",
    "            if 'EGFR' not in curr_mut:\n",
    "                non_egfr_count += 1\n",
    "            elif 'EGFR' in curr_mut:\n",
    "                egfr_count += 1\n",
    "                \n",
    "    msk_raw_drivers_subset.loc[idx].nonEGFR_muts = non_egfr_count\n",
    "    msk_raw_drivers_subset.loc[idx].EGFR_muts = egfr_count\n",
    "    \n",
    "    if non_egfr_count > 0 and egfr_count > 0 : # were def not seeing any samples with both \n",
    "        print(idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msk_raw_drivers_subset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set(msk_raw_drivers_subset.EGFR_muts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# want to binarize EGFR_muts values...dont want 'total mutational burden' to be skewed by the \n",
    "    # presence of multiple of the same EGFR muts\n",
    "def binarize(val):\n",
    "    if val > 0:\n",
    "        return(1)\n",
    "    else:\n",
    "        return(0)\n",
    "    \n",
    "egfr_muts_b = msk_raw_drivers_subset.EGFR_muts.apply(binarize)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add in standard and non-standard mutational burden columns\n",
    "msk_raw_drivers_subset['std_mut_burden'] = msk_raw_drivers_subset.EGFR_muts + msk_raw_drivers_subset.nonEGFR_muts\n",
    "msk_raw_drivers_subset['normalized_mut_burden'] = msk_raw_drivers_subset.nonEGFR_muts + egfr_muts_b\n",
    "msk_raw_drivers_subset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(set(msk_raw_drivers_subset.std_mut_burden))\n",
    "print(set(msk_raw_drivers_subset.normalized_mut_burden))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# high/low mutational burden...threshold = 1\n",
    "high_mut_burden_patients = list(msk_raw_drivers_subset[msk_raw_drivers_subset.std_mut_burden > 1].index)\n",
    "low_mut_burden_patients = list(msk_raw_drivers_subset[~msk_raw_drivers_subset.index.isin(high_mut_burden_patients)].index)\n",
    "\n",
    "print(len(high_mut_burden_patients))\n",
    "print(len(low_mut_burden_patients))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make some curves\n",
    "setup = kmf_setup_refined(meta_survival_sub, high_mut_burden_patients, low_mut_burden_patients)\n",
    "\n",
    "dur_in = setup[0]\n",
    "event_in = setup[1]\n",
    "dur_out = setup[2]\n",
    "event_out = setup[3]\n",
    "\n",
    "# create kmf object\n",
    "kmf = KaplanMeierFitter()\n",
    "\n",
    "# # plot\n",
    "kmf.fit(dur_in, event_in, label='high driver mutational burden')\n",
    "ax = kmf.plot(c='#cccc00', ci_show=False, show_censors=True, censor_styles={'ms': 9, 'marker': '|'}) # specify hex color code\n",
    "\n",
    "kmf.fit(dur_out, event_out, label='low driver mutational burden')\n",
    "ax = kmf.plot(ax=ax, c='#330066', ci_show=False, show_censors=True, censor_styles={'ms': 9, 'marker': '|'}) # specify hex color code\n",
    "\n",
    "# set axis labels\n",
    "ax.set_ylabel(\"surviving fraction\", fontname=\"Arial\", fontsize=12)\n",
    "ax.set_xlabel(\"timeline (years)\", fontname=\"Arial\", fontsize=12)\n",
    "\n",
    "# write to file\n",
    "fig = ax.get_figure()\n",
    "fig.savefig('../plot_out/NI12/breakdown_by_driver_mutational_burden.pdf')\n",
    "\n",
    "# get p value\n",
    "results = logrank_test(dur_in, dur_out, event_observed_A=event_in, event_observed_B=event_out)\n",
    "print(results.p_value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#///////////////////////////////////////////////////////////////////////////////////\n",
    "#///////////////////////////////////////////////////////////////////////////////////\n",
    "#//  lets look at 'overall' (tier1) mutational burden for our EGFR patients   //////\n",
    "#///////////////////////////////////////////////////////////////////////////////////"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# subset MSK tier1 dataframe by our EGFR+ LAUD patients\n",
    "print(len(msk_egfr_laud_patients))\n",
    "msk_egfr_sub_tier1 = msk[msk.index.isin(msk_egfr_laud_patients)]\n",
    "\n",
    "msk_egfr_sub_tier1['rowsum'] = msk_egfr_sub_tier1.sum(axis=1)\n",
    "msk_egfr_sub_tier1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set(msk_egfr_sub_tier1.rowsum)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "egfr_patients_hburden = list(msk_egfr_sub_tier1[msk_egfr_sub_tier1.rowsum > 0].index)\n",
    "egfr_patients_lburden = list(msk_egfr_sub_tier1[~msk_egfr_sub_tier1.index.isin(egfr_patients_hburden)].index)\n",
    "\n",
    "print(len(egfr_patients_hburden))\n",
    "print(len(egfr_patients_lburden))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make some curves\n",
    "setup = kmf_setup_refined(meta_survival_sub, egfr_patients_hburden, egfr_patients_lburden)\n",
    "\n",
    "dur_in = setup[0]\n",
    "event_in = setup[1]\n",
    "dur_out = setup[2]\n",
    "event_out = setup[3]\n",
    "\n",
    "# create kmf object\n",
    "kmf = KaplanMeierFitter()\n",
    "\n",
    "# # plot\n",
    "kmf.fit(dur_in, event_in, label='high mutational burden')\n",
    "ax = kmf.plot(c='#cccc00', ci_show=False, show_censors=True, censor_styles={'ms': 9, 'marker': '|'}) # specify hex color code\n",
    "\n",
    "kmf.fit(dur_out, event_out, label='low mutational ')\n",
    "ax = kmf.plot(ax=ax, c='#330066', ci_show=False, show_censors=True, censor_styles={'ms': 9, 'marker': '|'}) # specify hex color code\n",
    "\n",
    "# set axis labels\n",
    "ax.set_ylabel(\"surviving fraction\", fontname=\"Arial\", fontsize=12)\n",
    "ax.set_xlabel(\"timeline (years)\", fontname=\"Arial\", fontsize=12)\n",
    "\n",
    "# write to file\n",
    "fig = ax.get_figure()\n",
    "fig.savefig('../plot_out/NI12/egfr_patient_breakdown_low_high.pdf')\n",
    "\n",
    "# get p value\n",
    "results = logrank_test(dur_in, dur_out, event_observed_A=event_in, event_observed_B=event_out)\n",
    "print(results.p_value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "muts_to_drop = []\n",
    "for x in list(msk_egfr_sub_tier1.columns):\n",
    "    if 'EGFR' in x:\n",
    "        muts_to_drop.append(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get rid of the EGFR muts -- just want everything NOT egfr\n",
    "msk_egfr_sub_tier1 = msk_egfr_sub_tier1.drop(muts_to_drop, axis=1)\n",
    "msk_egfr_sub_tier1 = msk_egfr_sub_tier1.drop('rowsum', axis=1)\n",
    "msk_egfr_sub_tier1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msk_egfr_sub_tier1['rowsum'] = msk_egfr_sub_tier1.sum(axis=1)\n",
    "egfr_patients_hburden_no_egfr = list(msk_egfr_sub_tier1[msk_egfr_sub_tier1.rowsum > 0].index)\n",
    "egfr_patients_lburden_no_egfr = list(msk_egfr_sub_tier1[~msk_egfr_sub_tier1.index.isin(egfr_patients_hburden_no_egfr)].index)\n",
    "\n",
    "print(len(egfr_patients_hburden_no_egfr))\n",
    "print(len(egfr_patients_lburden_no_egfr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make some curves\n",
    "setup = kmf_setup_refined(meta_survival_sub, egfr_patients_hburden_no_egfr, egfr_patients_lburden_no_egfr)\n",
    "\n",
    "dur_in = setup[0]\n",
    "event_in = setup[1]\n",
    "dur_out = setup[2]\n",
    "event_out = setup[3]\n",
    "\n",
    "# create kmf object\n",
    "kmf = KaplanMeierFitter()\n",
    "\n",
    "# # plot\n",
    "kmf.fit(dur_in, event_in, label='high mutational burden patients')\n",
    "ax = kmf.plot(c='#cccc00', ci_show=False, show_censors=True, censor_styles={'ms': 9, 'marker': '|'}) # specify hex color code\n",
    "\n",
    "kmf.fit(dur_out, event_out, label='low mutational burden patients')\n",
    "ax = kmf.plot(ax=ax, c='#330066', ci_show=False, show_censors=True, censor_styles={'ms': 9, 'marker': '|'}) # specify hex color code\n",
    "\n",
    "# set axis labels\n",
    "ax.set_ylabel(\"surviving fraction\", fontname=\"Arial\", fontsize=12)\n",
    "ax.set_xlabel(\"timeline (years)\", fontname=\"Arial\", fontsize=12)\n",
    "\n",
    "# write to file\n",
    "fig = ax.get_figure()\n",
    "fig.savefig('../plot_out/NI12/egfr_patient_breakdown_low_high_cosmic_tier1.pdf')\n",
    "\n",
    "# get p value\n",
    "results = logrank_test(dur_in, dur_out, event_observed_A=event_in, event_observed_B=event_out)\n",
    "print(results.p_value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#///////////////////////////////////////////////////////////////////////////////////\n",
    "#///////////////////////////////////////////////////////////////////////////////////\n",
    "#//  lets look at 'overall' (total) mutational burden for our EGFR patients   //////\n",
    "#///////////////////////////////////////////////////////////////////////////////////"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this raw guy is the correct starting point for this analysis\n",
    "print(msk_raw_reorient_b.shape)\n",
    "\n",
    "# subset MSK raw dataframe by our EGFR+ LAUD patients\n",
    "print(len(msk_egfr_laud_patients))\n",
    "msk_egfr_sub_overall = msk_raw_reorient_b[msk_raw_reorient_b.index.isin(msk_egfr_laud_patients)]\n",
    "\n",
    "msk_egfr_sub_overall['rowsum'] = msk_egfr_sub_overall.sum(axis=1)\n",
    "msk_egfr_sub_overall.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set(msk_egfr_sub_overall.rowsum)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "high_burden_samples = list(msk_egfr_sub_overall[msk_egfr_sub_overall.rowsum > 2].index)\n",
    "low_burden_samples = list(msk_egfr_sub_overall[~msk_egfr_sub_overall.index.isin(high_burden_samples)].index)\n",
    "\n",
    "print(len(high_burden_samples))\n",
    "print(len(low_burden_samples))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make some curves -- this one looks really good\n",
    "setup = kmf_setup_refined(meta_survival_sub, high_burden_samples, low_burden_samples)\n",
    "\n",
    "dur_in = setup[0]\n",
    "event_in = setup[1]\n",
    "dur_out = setup[2]\n",
    "event_out = setup[3]\n",
    "\n",
    "# create kmf object\n",
    "kmf = KaplanMeierFitter()\n",
    "\n",
    "# # plot\n",
    "kmf.fit(dur_in, event_in, label='high mutational burden')\n",
    "ax = kmf.plot(c='#cccc00', ci_show=False, show_censors=True, censor_styles={'ms': 9, 'marker': '|'}) # specify hex color code\n",
    "\n",
    "kmf.fit(dur_out, event_out, label='low mutational burden')\n",
    "ax = kmf.plot(ax=ax, c='#330066', ci_show=False, show_censors=True, censor_styles={'ms': 9, 'marker': '|'}) # specify hex color code\n",
    "\n",
    "# set axis labels\n",
    "ax.set_ylabel(\"surviving fraction\", fontname=\"Arial\", fontsize=12)\n",
    "ax.set_xlabel(\"timeline (years)\", fontname=\"Arial\", fontsize=12)\n",
    "\n",
    "# write to file\n",
    "fig = ax.get_figure()\n",
    "fig.savefig('../plot_out/NI12/egfr_patient_breakdown_low_high_whole_exome.pdf')\n",
    "\n",
    "# get p value\n",
    "results = logrank_test(dur_in, dur_out, event_observed_A=event_in, event_observed_B=event_out)\n",
    "print(results.p_value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#////////////////////////////////////////////////////////////////////////\n",
    "#////////////////////////////////////////////////////////////////////////\n",
    "#////////////////////////////////////////////////////////////////////////"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
